from datetime import time
import numpy as np

from vnpy_ctastrategy import (
    CtaTemplate,
    StopOrder,
    TickData,
    BarData,
    TradeData,
    OrderData,
    ArrayManager,
)
from vnpy.trader.utility import BarGenerator, Interval


class MaStackWithMa20Slope(CtaTemplate):
    """
    入场：
        多头：pos=0 且 MA5 > MA10 且 MA5 > MA20 且（MA20回归斜率 a > slope_th 且可选R²过滤）
             （为避免重复开仓，采用“首次成立”判定：上一根未满足、当前根满足）
        空头：pos=0 且 MA5 < MA10 且 MA5 < MA20 且（MA20回归斜率 a < -slope_th 且可选R²过滤）
             （同样采用“首次成立”判定）

    出场：
        多头持仓（pos>0）：当 MA5 < MA10 且 MA5 < MA20 即离场
        空头持仓（pos<0）：当 MA5 > MA10 且 MA5 > MA20 即离场

    说明：
        - 仍使用 MA20 的线性回归斜率作为入场过滤，减少震荡误入。
        - 斜率：对最近 reg_k 个 MA20 做线性回归 y=a*x+b 的 a。
        - 可选 R² 过滤：use_r2=True 时要求 R² >= r2_min。
    """

    author: str = "用Python的交易员"

    # ====== 均线与下单参数 ======
    window_1: int = 5    # MA5
    window_2: int = 10   # MA10
    window_3: int = 20   # MA20
    add_price: int = 10  # 下单加价/滑点

    # ====== MA20回归斜率过滤参数 ======
    reg_k: int = 15
    slope_th: float = 0.0
    use_r2: bool = False
    r2_min: float = 0.2

    parameters = [
        "window_1", "window_2", "window_3",
        "add_price",
        "reg_k", "slope_th", "use_r2", "r2_min", "n_window",
    ]

    # 变量（面板可见）
    target: int = 0
    last_slope: float = 0.0
    last_r2: float = 0.0

    trail_up: float = 0.0
    trail_down: float = 0.0
    n_window: int = 5

    variables = [
        "target", "last_slope", "last_r2", "trail_up", "trail_down",
    ]

    def __init__(self, cta_engine, strategy_name, vt_symbol, setting) -> None:
        super().__init__(cta_engine, strategy_name, vt_symbol, setting)
        self.bg = BarGenerator(
            on_bar=None,
            on_window_bar=self.on_daily_bar,
            interval=Interval.DAILY,
            daily_end=time(14, 59),
        )
        self.am = ArrayManager()

    # ===== 生命周期 =====
    def on_init(self) -> None:
        self.write_log("策略初始化")
        init_bars = max(self.window_3, self.reg_k) + 5
        self.load_bar(init_bars)

    def on_start(self) -> None:
        self.write_log("策略启动")

    def on_stop(self) -> None:
        self.write_log("策略停止")

    # ===== 行情回调 =====
    def on_tick(self, tick: TickData) -> None:
        pass

    def on_bar(self, bar: BarData) -> None:
        if self.am.inited:
            if self.pos > 0:
                if bar.low_price <= self.trail_up - 2 * self.n:
                    self.sell(bar.close_price, 1, stop=True)
                    self.target = 0
            elif self.pos < 0:
                if bar.high_price >= self.trail_down + 2 * self.n:
                    self.cover(bar.close_price, 1, stop=True)
                    self.target = 0

        self.bg.update_bar(bar)

    def on_daily_bar(self, bar: BarData) -> None:
        # 1) 撤未成委托
        self.cancel_all()

        # 2) 更新数据
        self.am.update_bar(bar)
        if not self.am.inited:
            return

        # 3) 计算均线
        ma5  = self.am.sma(self.window_1, array=True)
        ma10 = self.am.sma(self.window_2, array=True)
        ma20 = self.am.sma(self.window_3, array=True)

        if len(ma5) < 2 or len(ma10) < 2 or len(ma20) < 2:
            return

        # 4) MA20 回归斜率（最近 reg_k 个 MA20）
        if len(ma20) >= self.reg_k:
            y = np.asarray(ma20[-self.reg_k:], dtype=float)
            x = np.arange(self.reg_k, dtype=float)
            a, b = np.polyfit(x, y, 1)   # a: slope
            slope_a = float(a)

            if self.use_r2:
                y_hat = a * x + b
                ss_res = float(np.sum((y - y_hat) ** 2))
                ss_tot = float(np.sum((y - np.mean(y)) ** 2))
                r2 = 0.0 if ss_tot == 0 else 1.0 - ss_res / ss_tot
            else:
                r2 = 1.0

            slope_up_ok = (slope_a >  self.slope_th) and (not self.use_r2 or r2 >= self.r2_min)
            slope_dn_ok = (slope_a < -self.slope_th) and (not self.use_r2 or r2 >= self.r2_min)
        else:
            slope_a, r2 = 0.0, 0.0
            slope_up_ok = False
            slope_dn_ok = False

        self.last_slope = slope_a
        self.last_r2 = r2

        # 5) “首次成立”判定：上一根不满足，本根满足（避免重复触发）
        up_now   = (ma5[-1] > ma10[-1] and ma5[-1] > ma20[-1])
        up_prev  = (ma5[-2] > ma10[-2] and ma5[-2] > ma20[-2])
        dn_now   = (ma5[-1] < ma10[-1] and ma5[-1] < ma20[-1])
        dn_prev  = (ma5[-2] < ma10[-2] and ma5[-2] < ma20[-2])

        # 6) 入场：MA5 相对 MA10/MA20 的“首次成立” × 斜率过滤
        if self.pos == 0:
            if up_now and not up_prev and slope_up_ok:
                self.buy(bar.close_price + self.add_price, 1)
                self.target = 1
            elif dn_now and not dn_prev and slope_dn_ok:
                self.short(bar.close_price - self.add_price, 1)
                self.target = -1
            else:
                self.target = 0

        # 7) 出场：仅按 MA5 相对 MA10/MA20 的当前关系
        else:
            # 多头：MA5 同时低于 MA10 和 MA20 → 离场
            if self.pos > 0 and (ma5[-1] < ma10[-1] and ma5[-1] < ma20[-1]):
                self.sell(bar.close_price, 1, stop=True)
                self.target = 0

            # 空头：MA5 同时高于 MA10 和 MA20 → 离场
            elif self.pos < 0 and (ma5[-1] > ma10[-1] and ma5[-1] > ma20[-1]):
                self.cover(bar.close_price, 1, stop=True)
                self.target = 0

        self.trail_up, self.trail_down = self.am.donchian(self.n_window)
        self.n = self.am.atr(self.n_window)

    # ===== 成交/委托回调 =====
    def on_trade(self, trade: TradeData) -> None:
        pass

    def on_order(self, order: OrderData) -> None:
        pass

    def on_stop_order(self, stop_order: StopOrder) -> None:
        pass
